import json
import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import pandas as pd
from tqdm import tqdm

from preprocessing import retinex


def rgb_to_msr(rgb_path, msr_root, config, mode='train'):
    if mode == 'train':
        video_path = os.path.join(msr_root, 'Train_files_MSRCR', Path(rgb_path).parent.name)
    elif mode == 'dev':
        video_path = os.path.join(msr_root, 'Dev_files_MSRCR', Path(rgb_path).parent.name)
    else:
        video_path = os.path.join(msr_root, 'Test_files_MSRCR', Path(rgb_path).parent.name)

    if not os.path.exists(video_path):
        try:
            os.mkdir(video_path)
        except OSError:
            pass
    msr_path = os.path.join(video_path, Path(rgb_path).name)

    if os.path.exists(os.path.join(video_path, Path(rgb_path).name)):
        return

    # 遍历该目录下的所有图片文件
    # print("filename :", rgb_path)
    img = cv2.imread(rgb_path)

    # change to gray
    # （下面第一行是将RGB转成单通道灰度图，第二步是将单通道灰度图转成3通道灰度图）
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # print("img.shape: ", img.shape)
    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    # print("img: ",img)
    img_MSRCR = retinex.MSRCR(
        img,
        config['sigma_list'],
        config['G'],
        config['b'],
        config['alpha'],
        config['beta'],
        config['low_clip'],
        config['high_clip']
    )
    # print("image_MSR.shape :", img_MSRCR.shape)
    # print("image_MSR :", img_MSRCR)



    # cv2.imwrite('D:\\Datasets\\test_2'+ "\\"+filename,img_MSRCR)
    cv2.imwrite(msr_path, img_MSRCR)


def get_images_paths(rgb_root, mode, data_csv):
    paths = []
    df = pd.read_csv(data_csv)
    for row in df.iterrows():
        if mode == 'train':
            path = os.path.join(rgb_root, 'Train_files_crops', row[1]['video'], row[1]['file'])
        elif mode == 'dev':
            path = os.path.join(rgb_root, 'Dev_files_crops', row[1]['video'], row[1]['file'])
        else:
            path = os.path.join(rgb_root, 'Test_files_crops', row[1]['video'], row[1]['file'])
        paths.append(path)
    return paths


def main():
    with open('config_msr.json', 'r') as f:
        config = json.load(f)
    # mode = 'train'
    mode = 'dev'
    # mode = 'test'
    rgb_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Oulu_NPU'
    msr_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Oulu_NPU'
    paths = get_images_paths(rgb_root, mode, '../data/data_{}.csv'.format(mode))
    with Pool(processes=2) as p:
        with tqdm(total=len(paths)) as pbar:
            for v in p.imap_unordered(partial(rgb_to_msr, msr_root=msr_root, config=config, mode=mode), paths):
                pbar.update()


if __name__ == "__main__":
    main()
